import json
import requests
from ai_configs import defalut_ai_config

current_ai_config = defalut_ai_config["doubao"]


def ai_chat(user_content, ai_config=current_ai_config):
    return ai_chat_system(user_content=user_content, ai_config=ai_config)


def ai_chat_system(user_content, prompt=None, ai_config=current_ai_config):
    message = [{"role": "user", "content": user_content}]
    if prompt:
        message.insert(0, {"role": "system", "content": prompt})
    return ai_chat_history(histories=message, ai_config=ai_config)


def ai_chat_history(histories=[], ai_config=current_ai_config):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {ai_config["key"]}",
    }
    response = requests.post(
        ai_config["url"],
        headers=headers,
        json={"model": ai_config["model"], "messages": histories},
    )
    if response.status_code == 200:
        response_txt = response.json()["choices"][0]["message"]["content"]
        # if isinstance(response_txt,str):

        return response_txt
    else:
        print(f"请求失败，状态码为{response.status_code}")


def default_on_receiving(full_txt, chunk_txt):
    print(chunk_txt, end="", flush=True)


# print(ai_chat("你好"))
def ai_chat_stream(
    ai_config,
    histories,
    on_think_start=None,
    on_think=None,
    on_think_end=None,
    on_receiving=None,
    on_finish=None,
):

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {ai_config["key"]}",
    }
    params = {"model": ai_config["model"], "messages": histories, "stream": True}
    try:
        with requests.post(
            url=ai_config["url"], headers=headers, json=params, stream=True, timeout=60
        ) as response:
            response.raise_for_status()
            response.encoding = "utf-8"
            think_full_text = ""
            receiving_full_text = ""
            is_think = False
            for line in response.iter_lines(decode_unicode=True):
                if line.startswith("data: ") and not line.startswith("data: [DONE]"):
                    line = json.loads(line[6:])
                    delta = line["choices"][0]["delta"]
                    # print(delta)
                    if "reasoning_content" in delta and bool(
                        delta["reasoning_content"]
                    ):
                        if not is_think:
                            is_think = True
                            if on_think_start is not None:
                                on_think_start()
                            think_full_text += delta["reasoning_content"]
                            if on_think is not None:
                                on_think(think_full_text, delta["reasoning_content"])
                            continue
                        if is_think:
                            think_full_text += delta["reasoning_content"]
                            if on_think is not None:
                                on_think(think_full_text, delta["reasoning_content"])
                            continue
                    if "content" in delta and bool(delta["content"]):
                        if is_think:
                            is_think = False
                            if on_think_end is not None:
                                on_think_end()
                            receiving_full_text += delta["content"]
                            if on_receiving is not None:
                                on_receiving(receiving_full_text, delta["content"])
                            continue
                        if not is_think:
                            receiving_full_text += delta["content"]
                            if on_receiving is not None:
                                on_receiving(receiving_full_text, delta["content"])
                            continue
            if on_finish is not None:
                on_finish(receiving_full_text)
    except Exception as e:
        print(f"发生异常:{e}")
